﻿#include "sampling.h"
#include "geometry.h"
/*
Original implementation from "A Low Distortion Map Between Disk and Square" by Peter Shirley et, al.
*
	Vector2 ToUnitDisk( Vector2 onSquare )
	real phi, r, u, v
	real a = 2*onSquare.X-1 // (a,b) is now on [-1,1]ˆ2
	real b = 2*onSquare.Y-1
	if (a > -b) // region 1 or 2
		if (a > b) // region 1, also |a| > |b|
			r = a
			phi = (PI/4 ) * (b/a)
		else // region 2, also |b| > |a|
			r = b;
			phi = (PI/4) * (2 - (a/b))
	else // region 3 or 4
	if (a < b) // region 3, also |a| >= |b|, a != 0
			r = -a
			phi = (PI/4) * (4 + (b/a))
		else // region 4, |b| >= |a|, but a==0 and b==0 could occur.
			r = -b
			if (b != 0)
			phi = (PI/4) * (6 - (a/b))
			else
			phi = 0
	u = r * cos( phi )
	v = r * sin( phi )
	return Vector2( u, v )
	end
*
*/
Point2 ConcentricSamplingDisk(const Point2 &sample) {
	// Map sample to [-1,1]
	float a = 2.f * sample.x - 1;
	float b = 2.f * sample.y - 1;
	Point2 coords;
	if(a == 0 && b == 0) {
		return Point2(0.0f,0.0f);
	}
	else if(a > -b) {
		if(a > b)
			coords = Point2(a, (M_PI/4.f) * b/a);
		else
			coords = Point2(b, (M_PI/4.f) * (2.f - a/b));
	} 
	else {
		if(a < b)
			coords = Point2(-a, (M_PI/4.f) * (4.f + b/a));
		else
			coords = Point2(-b, (M_PI/4.f) * (6.f - a/b));
	}
	return Point2(
		coords.x * std::cosf(coords.y),
		coords.y * std::sinf(coords.y)
		);
}

Point2 UniformSamplingDisk(const Point2 &sample)
{
	// Pdf = 1 / Area
	float r = sqrtf(sample.x);
	float theta = 2.f * M_PI * sample.y;
	return Point2(r*cosf(theta), r*sinf(theta));
}

Point2 RejectionSamplingDisk(RNG &rng) {
	// Pdf = 1 / Area
	float x,y;
	do {
		x = 1.f - 2.f * rng.RandomFloat();
		y = 1.f - 2.f * rng.RandomFloat();
	} while ((x*x + y*y) > 1.f);
	return Point2(x, y);
}

float UniformSamplingDiskPdf(float radius) {
	return INV_PI / (radius * radius);
}

Point2 UniformSamplingTriangle(const Point2 &sample) {
	// Pdf = 1 / Area
	float u = 1 - sqrtf(max(1-sample.x, 0.f));
	float v = (1 - u) * sample.y;
	return Point2(u, v);
}

float UniformSamplingTrianglePdf(float area) {
	return 1.f / area;
}

Vector3 UniformSamplingHemisphere(const Point2 &sample) {
	float z = sample.x;
	float r = sqrtf(max(0.f, 1.f - z*z));
	float phi = 2 * M_PI * sample.y;
	float x = r * cosf(phi);
	float y = r * sinf(phi);
	return Vector3(x, y, z);
}

float UniformSampingHemispherePdf() {
	return INV_TWOPI;
}

Vector3 UniformSamplingSphere(const Point2 &sample) {
	float z = 1 - 2 * sample.x;
	float r = sqrtf(max(0.f, 1.f - z*z));
	float phi = 2 * M_PI * sample.y;
	float x = r * cosf(phi);
	float y = r * sinf(phi);
	return Vector3(x, y, z);
}

float UniformSamplingSpherePdf() {
	return INV_FOURPI;
}

Vector3 CosineSamplingHemisphere(const Point2 &sample) {
	Point2 xy = ConcentricSamplingDisk(sample);
	float z = sqrtf(max(0.f, 1 - xy.x*xy.x - xy.y*xy.y));
	return Vector3(xy.x, xy.y, z);
}

float CosineSamplingHemispherePdf(float costheta) {
	return costheta * INV_PI;
}

Vector3 UniformSamplingCone(const Point2 &sample, float costhetamax) {
	float costheta = (1.f - sample.x) + sample.x * costhetamax;
	float sintheta = sqrtf(1.f - costheta*costheta);
	float phi = sample.y * 2.f * M_PI;
	return Vector3(cosf(phi) * sintheta, sinf(phi) * sintheta, costheta);
}

Vector3 UniformSamplingCone(const Point2 &sample, float costhetamax,
	const Vector3 &x, const Vector3 &y, const Vector3 &z) {
		float costheta = Lerp(sample.x, costhetamax, 1.f);
		float sintheta = sqrtf(1.f - costheta*costheta);
		float phi = sample.y * 2.f * M_PI;
		return cosf(phi) * sintheta * x + sinf(phi) * sintheta * y +
			costheta * z;
}
float UniformSamplingConePdf(float cosThetaMax) {
	return 1.f / (2.f * M_PI * (1.f - cosThetaMax));
}

void ComputeStep1dCDF(float *f, int nStep, float *c, float *cdf) {
	// Make sure:
	// f [0,nStep-1]
	// cdf[0,nStep];
	// c = sum(cdf);
	cdf[0] = 0;
	for(int i = 0; i < nStep + 1; i++)
		cdf[i] = cdf[i-1] + f[i-1]/nStep;
	*c = cdf[nStep];
	if(*c == 0.)	// make it uniform distribution
		for(int i = 0; i < nStep + 1; i++)
			cdf[i] = int(i)/nStep;	
	else 
		for(int i = 0; i < nStep + 1; i++)
		cdf[i] /= *c;
}

float SamplingStep1d(float *f, float *cdf, float c, int nStep, float u, float *pdf) {
	// A bug in pbrt v1 book P.644
	// lower_bound --> upper_bound or ptr-cdf-1 --> ptr-cdf
	// Here the return value is continues 
	// When sampling discrete, we can simply get (offset+u)/nStep * nStep
	// Discard the fractional, we get sample in position with offset
	float *ptr = std::upper_bound(cdf,cdf + nStep - 1, u);
	int offset = (int)(ptr-cdf-1);
	
	Assert(offset < nStep);
	Assert(u >= cdf[offset] && u < cdf[offset+1]);

	u = (u - cdf[offset]) / (cdf[offset+1] - cdf[offset]);
	
	Assert(!isnan(u));

	*pdf = f[offset] / c;
	return (offset + u ) / nStep;
}

void ComputeStep2dCFD(float *f, int nu, int nv, float *cdfu, float *cdfv, float *cu, float *cv) {
	// Make sure:
	// f [0, nu-1] x [0, nv-1]
	// cdfu [0, nu] x [0, nv-1]
	// cu [0, nv]
	// cdfv [0, nv];
	// cv [0, 0]
	//float *c = new float[nv]; // to store the sum cdf of each pConditionalV
	for (int i = 0; i < nv; i++)
	{
		// Compute conditional cdf of each v
		ComputeStep1dCDF(&f[i*nu], nu, &cu[i], &cdfu[i*(nu+1)]);
	}
	// Compute cdf for v, store cv 
	ComputeStep1dCDF(cu,nv,cv,cdfv);
}

Point2 SamplingStep2d(float *f, float *cdfu, float *cdfv, float *cu, float *cv, int nu, int nv, const Point2 &sample, float *pdf) {
	float fu,fv;
	float pdfu,pdfv;
	fv = SamplingStep1d(cu, cdfv, *cv, nv, sample.x, &pdfv);
	int iv = (int) (fv * nv);
	fu = SamplingStep1d(&f[iv*nu], &cdfu[iv*(nu+1)], cu[iv], nu, sample.y, &pdfu);
	*pdf = pdfu*pdfv;
	return Point2(fu,fv);
}

float SamplingStep2dPdf(float *f, float *cdfu, float *cdfv, float *cu, float *cv, int nu, int nv, float u, float v) {
	int iu = Clamp(Float2Int(u * nu), 0, nu-1);
	int iv = Clamp(Float2Int(v * nv), 0, nv-1);
	if(cu[iv] * (*cv) == 0.f) return 0.f;
	return (f[iv*nu+iu] * cu[iv]) / (cu[iv] * (*cv));
}

void StratifiedSample1D(float *samples, int n, bool jitter, RNG *rng) {
	float invN = 1.f / n;
	for(int i = 0; i < n; i++)
		samples[i] = min(OneMinusEpsilon, ((jitter ? rng->RandomFloat() : 0.5f) + i) * invN);
}

void StratifiedSample1D(vector<float> &samples, int n, bool jitter, RNG *rng) {
	samples.clear();
	samples.resize(n);
	float invN = 1.f / n;
	for(int i = 0; i < n; i++)
		samples[i] = min(OneMinusEpsilon, ((jitter ? rng->RandomFloat() : 0.5f) + i) * invN);
}

void StratifiedSample2D(Point2 *sample, int nx, int ny, bool jitter, RNG *rng) {
	float invX = 1.f / nx;
	float invY = 1.f / ny;
	int samp = 0;
	for(int j = 0; j < ny; j++)
		for(int i = 0; i < nx; i++) {
			sample[samp].x = min(OneMinusEpsilon, ((jitter ? rng->RandomFloat() : 0.5f) + i) * invX);
			sample[samp].y = min(OneMinusEpsilon, ((jitter ? rng->RandomFloat() : 0.5f) + j) * invY);
			samp++;
		}
}

void LatinHypercube1D(float *sample, int n, RNG *rng) {
	StratifiedSample1D(sample, n, true, rng);
	Shuffle(sample, n, rng);
}

void LatinHypercube2D(Point2 *sample, int n, RNG *rng) {
	float *x = new float[n];
	float *y = new float[n];
	
	StratifiedSample1D(x, n, true, rng);
	StratifiedSample1D(y, n, true, rng);
	
	Shuffle(y, n, rng);
	Shuffle(x, n, rng);
	
	for(int i = 0; i < n; i++)
		sample[i] = Point2(x[i], y[i]);
	
	SAFE_DELETE_ARRAY(x);
	SAFE_DELETE_ARRAY(y);
}

PermutedHalton::PermutedHalton(uint32_t d) {
	rng = new RNG(time(NULL));
	dims = d;
	// Determine bases $b_i$ and their sum
	b = new uint32_t[dims];
	uint32_t sumBases = 0;
	for (uint32_t i = 0; i < dims; ++i) {
		b[i] = primes[i];
		sumBases += b[i];
	}

	// Compute permutation tables for each base
	permute = new uint32_t[sumBases];
	uint32_t *p = permute;
	for (uint32_t i = 0; i < dims; ++i) {
		GeneratePermutation(p, b[i], rng);
		p += b[i];
	}
}

/*
	Poisson Disk Sampling Helper Functions
*/

// sample annulus
void sampleAnnulus(float radius, const Vector2 &center, RNG *rng, Vector2 *sample) {
	float u = (rng->RandomFloat() + 1) * radius;		// u = [r, 2r];
	float v = rng->RandomFloat() * 2 * M_PI;			// v = [0, 2PI]
	sample->x = center.x + u * cosf(v);
	sample->y = center.y + u * sinf(v);
}
// locate the index of a sample
int arrayIndex(Vector2 dx, Vector2 dimensions, Vector2 sample) {
	int i = Float2Int(sample.x / dx.x);
	int j = Float2Int(sample.y / dx.y);
	return i + j * (int)dimensions.y;
}
// 2D blue noise sampling for one grid pixel
void bluenoiseSampling(float radius, vector<Vector2> &samples, RNG *rng, int maxAttemplts) {
	
	samples.clear();
	vector<uint> activeList;

	// acceleration grid

	float gridDx = 0.999f * radius/sqrtf(2);	// a grid cell this size can have at most one sample in it
	Vector2 dx(gridDx, gridDx);
	Vector2 dimensions;
	dimensions.x = dimensions.y = Ceil2Int(1.f / gridDx);
	uint32_t totalArraySize = dimensions.x * dimensions.y;

	vector<int> accel(totalArraySize, -1);	// -1 indicates no sample, otherwise the index of sample

	// first sample
	Vector2 x;
	x.x = rng->RandomFloat(); x.y = rng->RandomFloat();
	
	samples.push_back(x);
	activeList.push_back(0);
	uint k = arrayIndex(dx, dimensions, x);
	accel[k] = 0;

	while(!activeList.empty()) {
		uint r = int ((activeList.size() - 1) * rng->RandomFloat());
		int p = activeList[r];
		bool foundSample = false;
		Vector2 j, jmin, jmax;
		for(int attempt = 0; attempt < maxAttemplts; attempt++) {
			sampleAnnulus(radius, samples[p], rng, &x);
			// check this sample is within bounds
			if(x.x < 0. || x.x > 1.f || x.y < 0. || x.y > 1.)
				goto rejectSample;
			// test proximity to nearby samples
			for(int i = 0; i < 2; i++) {
				int thismin = (int) ((x[i] - radius) / gridDx);
				if(thismin < 0)
					thismin = 0;
				else if(thismin >= (int) dimensions[i])
					thismin = dimensions[i] - 1;
				jmin[i] = thismin;
				int thismax = (int) ((x[i] + radius) / gridDx);
				if(thismax < 0)
					thismax = 0;
				if(thismax >= (int) dimensions[i])
					thismax = dimensions[i] - 1;
				jmax[i] = thismax;
			}
			for(j = jmin;;) {
				// check if there's a sample at j that's too close to x
				k = arrayIndex(dx, dimensions, j * gridDx);
				if(accel[k] >= 0 && accel[k] != p) {
					if(DistanceSqured(x, samples[accel[k]]) < radius * radius)
						goto rejectSample;
				}
				// move to next j
				for(uint i = 0; i < 2; i++) {
					++j[i];
					if(j[i] <= jmax[i])
						break;
					else {
						if( i == 1)
							goto doneJLoop;
						else
							j[i] = jmin[i];
					}
				}
			}
			doneJLoop:
			// if we made it here, we're good
			foundSample = true;
			break;
			rejectSample:
			; // nothing to do except the next iteration
		}
		if(foundSample) {
			size_t q = samples.size();
			samples.push_back(x);
			activeList.push_back(q);
			k = arrayIndex(dx, dimensions, x);
			accel[k] = (int)q;
		}
		else {
			//since we couldn't find a sample on p's disk, we remove p from the active list
			activeList[r] = activeList.back();
			activeList.pop_back();
		}
	}

}